# core/query/query_refiner.py
import pandas as pd
from core.model.model_caller import ModelCaller
from core.utils.logging_handler import LoggingHandler

class QueryRefiner:
    def __init__(self, model_caller, df=None):
        self.model_caller = model_caller
        self.df = df
        self.logger = LoggingHandler().get_logger()

    def _get_column_value_info(self):
        """生成每列数据的描述信息"""
        column_info = []
        for column in self.df.columns:
            if pd.api.types.is_numeric_dtype(self.df[column]):
                col_min = self.df[column].min()
                col_max = self.df[column].max()
                column_info.append(f"{column} (数值): 最小值 = {col_min}, 最大值 = {col_max}")
            elif pd.api.types.is_categorical_dtype(self.df[column]) or self.df[column].dtype == object:
                unique_values = self.df[column].unique()
                column_info.append(f"{column} (类别): 可能的值 = {', '.join(map(str, unique_values))}")
            else:
                column_info.append(f"{column} (其他类型)")
        return "\n".join(column_info)

    def refine_query_with_model(self, user_input, use_remote=True):
        self.logger.info("开始精确化用户查询")

        # 获取列名和列的取值范围信息
        df_columns_info = self._get_column_value_info()

        prompt = (
            "背景信息：\n"
            "- 用户的问题可能是与数据表中的信息相关的查询，或者与数据表无关的闲聊。\n"
            "- 数据表包括多个品牌和型号的电气产品，以及这些产品的各项性能指标。\n"
            "- 对于查询请求，生成的代码和结果必须基于数据表中的信息，确保回答具体且相关。\n"
            "- 对于闲聊请求，不需要生成查询代码，而是直接用自然语言回答用户。\n\n"
            "注意事项：\n"
            "- 用户提问中的品牌名和产品名可能包含中英文混合或缩写，请根据数据表中提供的取值范围进行匹配。\n"
            "- 在生成查询时，优先选择与数据表中列名和取值范围最匹配的品牌名、产品名及其他相关信息。\n"
            "- 生成查询代码时，必须确保包含明确的精确查询指令，并使用数据表中的具体列名和取值。\n"
            "- 如果用户提出的是闲聊问题，模型的输出中必须包含 '闲聊应答'，并提供自然语言的回复内容。\n\n"
            "表格 'df' 包含以下列名和取值范围：\n\n"
            f"{df_columns_info}\n\n"
            "任务：\n"
            f"用户输入了一个请求：'{user_input}'。\n"
            "请根据以下条件处理用户请求：\n"
            "- 如果用户的请求是与数据表相关的查询问题，请将其转化为一个更为精确、具体的查询。\n"
            "  - 查询应明确地表达用户的意图，并使用数据表中的具体列名和取值。\n"
            "  - 基于数据表中的列名和取值范围，指定要查询的具体列以及相关的过滤条件或聚合操作。\n"
            "  - 生成的查询必须包含明确的精确查询指令。\n"
            "- 如果用户的请求是与数据表无关的闲聊问题，请直接用自然语言回答用户的闲聊内容，输出中必须包含 '闲聊应答'，并提供相应的回复。\n\n"
            "可以参考以下示例来帮助构建精确的查询或应对闲聊：\n\n"
            "示例1: \n"
            "用户输入: '施耐德和ABB哪个品牌更好？'\n"
            "精确查询: '请比较Brand列中施耐德（Schneider）和ABB两者在Electrical durability AC3 400V (M cycles)列上的平均值，来判断哪个品牌的AC3电寿命更优。'\n\n"
            "示例2: \n"
            "用户输入: 'AC控制和DC控制哪个更省电？'\n"
            "精确查询: '请比较Coil Control列中AC和DC两种控制方式在Coil Holding Power(W/VA)列上的平均值，来判断哪种控制方式更省电。'\n\n"
            "示例3: \n"
            "用户输入: '比较一下不同系列的电气寿命。'\n"
            "精确查询: '请按Series列分组，并展示Electrical durability AC3 400V (M cycles)列的平均值，来比较不同系列的电气寿命。'\n\n"
            "示例4: \n"
            "用户输入: '今天的天气怎么样？'\n"
            "闲聊应答: '闲聊应答：今天的天气预报显示，气温在20到25度之间，建议带把伞以防万一。'\n\n"
            "请基于以上信息，生成一个适合用户需求的响应。如果是查询请求，请生成查询代码；如果是闲聊，请直接回复闲聊内容，并在输出中包含 '闲聊应答'。"
        )


        self.logger.info(f"生成的提示词: {prompt}")

        # 调用封装的模型接口，选择远程或本地模型
        refined_query = self.model_caller.call_model(prompt, use_remote=use_remote)
        self.logger.info(f"生成的精确查询: {refined_query}")
        return refined_query

